import logging
import socket
import sys
from Crypto.Cipher import ARC4

import rsa
from rsa import PublicKey

sys.path.append('..')
from transfer.common import randomword, get_file_sha1
from transfer import sender

PublicKey

LOG = logging.getLogger(__name__)
logging.basicConfig(level=logging.DEBUG)

SERVER_HOST = socket.gethostname()

PUBLIC_KEY, PRIVATE_KEY = rsa.newkeys(1024)
PUBLIC_KEY_LEN = len(str(PUBLIC_KEY)) + 1
RC4_KEY_LEN = 128


class Receiver(object):
    def __init__(self, host, port):
        self.s = socket.socket()
        self.sender_public_key = None
        self.rc4_key = None
        self.server_host = host
        self.port = port
        self.connect()

    def connect(self):
        self.s.connect((self.server_host, self.port))

    def close(self):
        self.s.close()

    def shake_hand(self):
        LOG.debug('shaking hand with %s', str((SERVER_HOST, self.port)))
        self.sender_public_key = eval(self.s.recv(PUBLIC_KEY_LEN))
        LOG.debug('receive public key: %s' % self.sender_public_key)
        self.s.send(str(PUBLIC_KEY))
        msg = self.s.recv(1024)
        self.rc4_key = rsa.decrypt(msg, PRIVATE_KEY)
        msg = randomword(RC4_KEY_LEN / 2)
        self.rc4_key += msg
        self.s.send(rsa.encrypt(msg, self.sender_public_key))
        LOG.debug('rc4 key: %s' % self.rc4_key)

    def receive_file(self, file_path):
        with open('%s' % file_path, 'wb') as f:
            LOG.debug('open file:%s' % file_path)
            cipher = ARC4.new(self.rc4_key)
            while True:
                data = self.s.recv(1024)
                if not data:
                    break
                plain = cipher.decrypt(data)
                f.write(plain)

        f.close()
        self.close()
        LOG.debug('Successfully get the file')

    def check_file(self, file_path):
        self.shake_hand()
        signature = self.s.recv(1024)
        LOG.debug(len(signature))
        message = get_file_sha1('%s' % file_path)
        done = rsa.verify(message, signature, self.sender_public_key)
        if not done:
            raise Exception('file destroy')
        LOG.debug(done)
        self.close()
        return done


if __name__ == '__main__':
    client = Receiver(sender.LOCALHOST, sender.PORT_FILE)
    client.shake_hand()
    client.receive_file('out.pdf')
    client1 = Receiver(sender.LOCALHOST, sender.PORT_DIGEST)
    client1.check_file('out.pdf')
